#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PgBouncer管理模块
提供PgBouncer连接池的安装、配置及管理功能
"""

import os
import re
import shutil
import psutil
import platform
import subprocess
from typing import Dict, List, Optional, Tuple, Any

from ..core.utils import (
    logger, run_command, is_admin, is_windows, 
    command_exists, confirm_action, get_input_with_default
)
from ..core.ui_utils import (
    print_title, print_info, print_success, 
    print_error, print_warning, get_input, confirm
)

class PgBouncerManager:
    """PgBouncer管理类"""
    
    def __init__(self):
        """初始化PgBouncer管理器"""
        self.service_name = "pgbouncer"
        self.config_dir = "/etc/pgbouncer" if not is_windows() else "C:\\pgbouncer"
        self.config_file = os.path.join(self.config_dir, "pgbouncer.ini")
        self.userlist_file = os.path.join(self.config_dir, "userlist.txt")
        
    def install(self):
        """安装PgBouncer"""
        print_title("安装PgBouncer")
        
        # 检查是否已安装
        if self._is_installed():
            if not confirm("PgBouncer已安装，是否重新安装?"):
                return
                
        # 检查管理员权限
        if not is_admin():
            print_error("安装PgBouncer需要管理员权限，请以管理员身份运行")
            return
            
        try:
            if is_windows():
                self._install_windows()
            else:
                self._install_linux()
        except Exception as e:
            print_error(f"安装PgBouncer时出错: {e}")
            logger.exception("安装PgBouncer异常")
    
    def _install_windows(self):
        """在Windows上安装PgBouncer"""
        print_info("Windows系统安装PgBouncer...")
        print_error("Windows平台暂不支持自动安装PgBouncer")
        print_info("请按照以下步骤手动安装:")
        print_info("1. 从 https://www.pgbouncer.org/downloads/ 下载Windows版本")
        print_info("2. 解压到C:\\pgbouncer目录")
        print_info("3. 创建配置文件C:\\pgbouncer\\pgbouncer.ini")
        print_info("4. 使用NSSM工具将PgBouncer注册为Windows服务")
    
    def _install_linux(self):
        """在Linux上安装PgBouncer"""
        print_info("Linux系统安装PgBouncer...")
        
        # 检测Linux发行版
        if os.path.exists("/etc/debian_version"):
            # Debian/Ubuntu系统
            print_info("检测到Debian/Ubuntu系统")
            run_command(["apt-get", "update"])
            run_command(["apt-get", "install", "-y", "pgbouncer"])
        elif os.path.exists("/etc/redhat-release"):
            # RHEL/CentOS系统
            print_info("检测到RHEL/CentOS系统")
            run_command(["yum", "install", "-y", "pgbouncer"])
        else:
            print_warning("未知的Linux发行版，尝试使用apt安装")
            try:
                run_command(["apt-get", "update"])
                run_command(["apt-get", "install", "-y", "pgbouncer"])
            except Exception as e:
                print_error(f"安装失败: {e}")
                print_info("请尝试手动安装PgBouncer")
                return
        
        print_success("PgBouncer安装成功")
        
        # 检查pgbouncer配置目录是否存在
        if not os.path.exists(self.config_dir):
            print_warning(f"配置目录 {self.config_dir} 不存在，创建中...")
            try:
                os.makedirs(self.config_dir, exist_ok=True)
            except Exception as e:
                print_error(f"创建配置目录失败: {e}")
                return
        
        # 如果初始配置文件不存在，创建一个简单的模板
        if not os.path.exists(self.config_file):
            print_info("初始配置文件不存在，创建基础配置...")
            self._create_default_config()
        
        # 创建基本用户列表文件
        if not os.path.exists(self.userlist_file):
            print_info("用户列表文件不存在，创建空白用户列表...")
            self._create_empty_userlist()
            
        # 启动服务
        print_info("启动PgBouncer服务...")
        self.start()
    
    def _create_default_config(self):
        """创建默认的PgBouncer配置文件"""
        try:
            with open(self.config_file, 'w') as f:
                f.write("""[databases]
* = host=localhost port=5432

[pgbouncer]
logfile = /var/log/postgresql/pgbouncer.log
pidfile = /var/run/postgresql/pgbouncer.pid
listen_addr = *
listen_port = 6432
auth_type = md5
auth_file = /etc/pgbouncer/userlist.txt
admin_users = postgres
stats_users = postgres
pool_mode = session
server_reset_query = DISCARD ALL
max_client_conn = 100
default_pool_size = 20
min_pool_size = 10
reserve_pool_size = 5
reserve_pool_timeout = 3
max_db_connections = 0
max_user_connections = 0
tcp_keepalive = 1
tcp_keepcnt = 5
tcp_keepidle = 30
tcp_keepintvl = 30
""")
            print_success("创建默认配置文件成功")
        except Exception as e:
            print_error(f"创建配置文件失败: {e}")
    
    def _create_empty_userlist(self):
        """创建空的用户列表文件"""
        try:
            with open(self.userlist_file, 'w') as f:
                f.write('"postgres" ""\n')
            print_success("创建用户列表文件成功")
        except Exception as e:
            print_error(f"创建用户列表文件失败: {e}")
    
    def _is_installed(self) -> bool:
        """检查PgBouncer是否已安装"""
        if is_windows():
            # Windows上检查程序目录
            return os.path.exists("C:\\pgbouncer\\pgbouncer.exe")
        else:
            # Linux上检查命令是否存在
            return command_exists("pgbouncer")
    
    def uninstall(self):
        """卸载PgBouncer"""
        print_title("卸载PgBouncer")
        
        # 检查是否已安装
        if not self._is_installed():
            print_error("PgBouncer未安装，无需卸载")
            return
            
        # 检查管理员权限
        if not is_admin():
            print_error("卸载PgBouncer需要管理员权限，请以管理员身份运行")
            return
            
        # 确认卸载
        if not confirm("确定要卸载PgBouncer吗?"):
            return
            
        # 停止服务
        print_info("停止PgBouncer服务...")
        self.stop()
        
        try:
            if is_windows():
                self._uninstall_windows()
            else:
                self._uninstall_linux()
        except Exception as e:
            print_error(f"卸载PgBouncer时出错: {e}")
            logger.exception("卸载PgBouncer异常")
    
    def _uninstall_windows(self):
        """在Windows上卸载PgBouncer"""
        print_info("Windows系统卸载PgBouncer...")
        print_warning("Windows平台暂不支持自动卸载PgBouncer")
        print_info("请按照以下步骤手动卸载:")
        print_info("1. 使用NSSM工具移除PgBouncer服务")
        print_info("2. 删除C:\\pgbouncer目录")
    
    def _uninstall_linux(self):
        """在Linux上卸载PgBouncer"""
        print_info("Linux系统卸载PgBouncer...")
        
        # 检测Linux发行版
        if os.path.exists("/etc/debian_version"):
            # Debian/Ubuntu系统
            print_info("检测到Debian/Ubuntu系统")
            run_command(["apt-get", "remove", "-y", "pgbouncer"])
            if confirm("是否同时删除配置文件?"):
                run_command(["apt-get", "purge", "-y", "pgbouncer"])
        elif os.path.exists("/etc/redhat-release"):
            # RHEL/CentOS系统
            print_info("检测到RHEL/CentOS系统")
            run_command(["yum", "remove", "-y", "pgbouncer"])
        else:
            print_warning("未知的Linux发行版，尝试使用apt卸载")
            try:
                run_command(["apt-get", "remove", "-y", "pgbouncer"])
            except Exception as e:
                print_error(f"卸载失败: {e}")
                print_info("请尝试手动卸载PgBouncer")
                return
        
        print_success("PgBouncer卸载成功")
    
    def start(self):
        """启动PgBouncer服务"""
        print_info("启动PgBouncer服务...")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if is_windows():
            print_warning("Windows平台请手动启动PgBouncer服务")
            return
            
        try:
            run_command(["sudo", "systemctl", "start", self.service_name])
            print_success("PgBouncer服务启动成功")
        except Exception as e:
            print_error(f"启动PgBouncer服务失败: {e}")
    
    def stop(self):
        """停止PgBouncer服务"""
        print_info("停止PgBouncer服务...")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，无需停止")
            return
            
        if is_windows():
            print_warning("Windows平台请手动停止PgBouncer服务")
            return
            
        try:
            run_command(["sudo", "systemctl", "stop", self.service_name])
            print_success("PgBouncer服务停止成功")
        except Exception as e:
            print_error(f"停止PgBouncer服务失败: {e}")
    
    def restart(self):
        """重启PgBouncer服务"""
        print_info("重启PgBouncer服务...")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if is_windows():
            print_warning("Windows平台请手动重启PgBouncer服务")
            return
            
        try:
            run_command(["sudo", "systemctl", "restart", self.service_name])
            print_success("PgBouncer服务重启成功")
        except Exception as e:
            print_error(f"重启PgBouncer服务失败: {e}")
    
    def reload(self):
        """重载PgBouncer配置"""
        print_info("重载PgBouncer配置...")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if is_windows():
            print_warning("Windows平台请手动重载PgBouncer配置")
            return
            
        try:
            run_command(["sudo", "systemctl", "reload", self.service_name])
            print_success("PgBouncer配置重载成功")
        except Exception as e:
            print_error(f"重载PgBouncer配置失败: {e}")
            print_info("尝试使用信号方式重载...")
            try:
                run_command(["sudo", "pkill", "-SIGHUP", "pgbouncer"])
                print_success("通过信号方式重载PgBouncer配置成功")
            except Exception as e:
                print_error(f"通过信号重载配置失败: {e}")
    
    def status(self):
        """查看PgBouncer服务状态"""
        print_info("查看PgBouncer服务状态...")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if is_windows():
            print_warning("Windows平台请手动查看PgBouncer服务状态")
            return
            
        try:
            run_command(["sudo", "systemctl", "status", self.service_name])
        except Exception as e:
            print_error(f"查看PgBouncer服务状态失败: {e}")
    
    def show_config(self):
        """显示PgBouncer配置文件路径"""
        print_info("PgBouncer配置文件:")
        print_info(f"主配置文件: {self.config_file}")
        print_info(f"用户列表文件: {self.userlist_file}")
        
        if os.path.exists(self.config_file):
            if confirm("是否查看配置文件内容?"):
                try:
                    with open(self.config_file, 'r') as f:
                        print(f.read())
                except Exception as e:
                    print_error(f"读取配置文件失败: {e}")
        else:
            print_warning(f"配置文件 {self.config_file} 不存在")
    
    def optimize_config(self):
        """根据服务器硬件优化PgBouncer配置"""
        print_title("优化PgBouncer配置")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if not os.path.exists(self.config_file):
            print_error(f"配置文件 {self.config_file} 不存在，请先安装PgBouncer")
            return
            
        print_info("正在分析系统资源...")
        
        # 获取系统信息
        memory = psutil.virtual_memory()
        total_memory_mb = memory.total / (1024 * 1024)  # 转换为MB
        cpu_count = psutil.cpu_count(logical=True)
        
        print_info(f"系统总内存: {total_memory_mb:.0f} MB")
        print_info(f"CPU核心数: {cpu_count}")
        
        # 计算优化参数
        # 每个PgBouncer连接大约占用2KB内存，但PostgreSQL后端连接消耗更多
        # 假设PostgreSQL连接平均占用约10MB内存
        # 预留40%系统内存给操作系统和其他应用
        available_memory = total_memory_mb * 0.6
        
        # 计算最大客户端连接数 - 取决于应用需求和系统资源
        # 避免设置过高以防系统打开文件描述符不足
        max_conn = min(int(available_memory / 10), 5000)  # 限制最大值为5000
        
        # 默认连接池大小 - 根据CPU核心数和预期负载调整
        # 每个池中的连接太少会导致频繁创建/销毁连接，太多则浪费资源
        default_pool_size = max(20, min(int(cpu_count * 5), 50))
        
        # 最小连接池大小 - 保持一定数量的空闲连接以应对突发负载
        min_pool_size = max(5, int(default_pool_size / 4))
        
        # 保留连接池大小 - 用于应对临时高负载
        reserve_pool_size = max(5, int(default_pool_size / 2))
        
        # 其他参数
        tcp_keepalive = 1
        tcp_keepidle = 30
        tcp_keepintvl = 15
        tcp_keepcnt = 6
        
        # 确认是否应用优化
        print_info("\n计算得出的优化参数:")
        print_info(f"最大客户端连接数 (max_client_conn): {max_conn}")
        print_info(f"默认连接池大小 (default_pool_size): {default_pool_size}")
        print_info(f"最小连接池大小 (min_pool_size): {min_pool_size}")
        print_info(f"保留连接池大小 (reserve_pool_size): {reserve_pool_size}")
        print_info(f"TCP保活时间 (tcp_keepidle): {tcp_keepidle}秒")
        print_info(f"TCP保活间隔 (tcp_keepintvl): {tcp_keepintvl}秒")
        print_info(f"TCP保活探测次数 (tcp_keepcnt): {tcp_keepcnt}次")
        print_info(f"监听地址 (listen_addr): * (所有接口)")
        
        if not confirm("是否应用这些优化参数?"):
            return
            
        # 备份当前配置
        backup_file = f"{self.config_file}.bak"
        try:
            shutil.copy2(self.config_file, backup_file)
            print_info(f"已备份原配置到 {backup_file}")
        except Exception as e:
            print_error(f"备份配置文件失败: {e}")
            if not confirm("继续而不备份?"):
                return
        
        # 读取当前配置
        try:
            with open(self.config_file, 'r') as f:
                config_lines = f.readlines()
        except Exception as e:
            print_error(f"读取配置文件失败: {e}")
            return
        
        # 更新配置参数
        new_config = []
        params_updated = {
            "max_client_conn": False,
            "default_pool_size": False,
            "min_pool_size": False,
            "reserve_pool_size": False,
            "tcp_keepalive": False,
            "tcp_keepidle": False,
            "tcp_keepintvl": False,
            "tcp_keepcnt": False,
            "pool_mode": False,
            "server_reset_query": False,
            "server_check_delay": False,
            "server_check_query": False,
            "server_lifetime": False,
            "server_idle_timeout": False,
            "query_wait_timeout": False,
            "listen_addr": False
        }
        
        in_pgbouncer_section = False
        for line in config_lines:
            # 识别pgbouncer部分
            if line.strip() == "[pgbouncer]":
                in_pgbouncer_section = True
                new_config.append(line)
                continue
            
            # 处理新部分开始
            if line.strip().startswith("[") and line.strip() != "[pgbouncer]":
                in_pgbouncer_section = False
                
            # 在pgbouncer部分内更新参数
            if in_pgbouncer_section:
                # 检查并更新参数
                if line.strip().startswith("max_client_conn"):
                    new_config.append(f"max_client_conn = {max_conn}\n")
                    params_updated["max_client_conn"] = True
                elif line.strip().startswith("default_pool_size"):
                    new_config.append(f"default_pool_size = {default_pool_size}\n")
                    params_updated["default_pool_size"] = True
                elif line.strip().startswith("min_pool_size"):
                    new_config.append(f"min_pool_size = {min_pool_size}\n")
                    params_updated["min_pool_size"] = True
                elif line.strip().startswith("reserve_pool_size"):
                    new_config.append(f"reserve_pool_size = {reserve_pool_size}\n")
                    params_updated["reserve_pool_size"] = True
                elif line.strip().startswith("tcp_keepalive"):
                    new_config.append(f"tcp_keepalive = {tcp_keepalive}\n")
                    params_updated["tcp_keepalive"] = True
                elif line.strip().startswith("tcp_keepidle"):
                    new_config.append(f"tcp_keepidle = {tcp_keepidle}\n")
                    params_updated["tcp_keepidle"] = True
                elif line.strip().startswith("tcp_keepintvl"):
                    new_config.append(f"tcp_keepintvl = {tcp_keepintvl}\n")
                    params_updated["tcp_keepintvl"] = True
                elif line.strip().startswith("tcp_keepcnt"):
                    new_config.append(f"tcp_keepcnt = {tcp_keepcnt}\n")
                    params_updated["tcp_keepcnt"] = True
                elif line.strip().startswith("pool_mode"):
                    new_config.append("pool_mode = session\n")
                    params_updated["pool_mode"] = True
                elif line.strip().startswith("server_reset_query"):
                    new_config.append("server_reset_query = DISCARD ALL\n")
                    params_updated["server_reset_query"] = True
                elif line.strip().startswith("server_check_delay"):
                    new_config.append("server_check_delay = 30\n")
                    params_updated["server_check_delay"] = True
                elif line.strip().startswith("server_check_query"):
                    new_config.append("server_check_query = SELECT 1\n")
                    params_updated["server_check_query"] = True
                elif line.strip().startswith("server_lifetime"):
                    # 设置服务器连接生命周期为1小时
                    new_config.append("server_lifetime = 3600\n")
                    params_updated["server_lifetime"] = True
                elif line.strip().startswith("server_idle_timeout"):
                    # 设置服务器空闲超时为10分钟
                    new_config.append("server_idle_timeout = 600\n")
                    params_updated["server_idle_timeout"] = True
                elif line.strip().startswith("query_wait_timeout"):
                    # 设置查询等待超时为120秒
                    new_config.append("query_wait_timeout = 120\n")
                    params_updated["query_wait_timeout"] = True
                elif line.strip().startswith("listen_addr"):
                    # 设置监听地址为所有接口
                    new_config.append("listen_addr = *\n")
                    params_updated["listen_addr"] = True
                else:
                    new_config.append(line)
            else:
                new_config.append(line)
        
        # 检查是否完成了pgbouncer部分的更新
        if in_pgbouncer_section:
            # 添加缺少的参数
            for param, updated in params_updated.items():
                if not updated:
                    if param == "max_client_conn":
                        new_config.append(f"max_client_conn = {max_conn}\n")
                    elif param == "default_pool_size":
                        new_config.append(f"default_pool_size = {default_pool_size}\n")
                    elif param == "min_pool_size":
                        new_config.append(f"min_pool_size = {min_pool_size}\n")
                    elif param == "reserve_pool_size":
                        new_config.append(f"reserve_pool_size = {reserve_pool_size}\n")
                    elif param == "tcp_keepalive":
                        new_config.append(f"tcp_keepalive = {tcp_keepalive}\n")
                    elif param == "tcp_keepidle":
                        new_config.append(f"tcp_keepidle = {tcp_keepidle}\n")
                    elif param == "tcp_keepintvl":
                        new_config.append(f"tcp_keepintvl = {tcp_keepintvl}\n")
                    elif param == "tcp_keepcnt":
                        new_config.append(f"tcp_keepcnt = {tcp_keepcnt}\n")
                    elif param == "pool_mode":
                        new_config.append("pool_mode = session\n")
                    elif param == "server_reset_query":
                        new_config.append("server_reset_query = DISCARD ALL\n")
                    elif param == "server_check_delay":
                        new_config.append("server_check_delay = 30\n")
                    elif param == "server_check_query":
                        new_config.append("server_check_query = SELECT 1\n")
                    elif param == "server_lifetime":
                        new_config.append("server_lifetime = 3600\n")
                    elif param == "server_idle_timeout":
                        new_config.append("server_idle_timeout = 600\n")
                    elif param == "query_wait_timeout":
                        new_config.append("query_wait_timeout = 120\n")
                    elif param == "listen_addr":
                        new_config.append("listen_addr = *\n")
        
        # 如果没有找到pgbouncer部分，则添加
        if "[pgbouncer]" not in "".join(config_lines):
            new_config.append("\n[pgbouncer]\n")
            new_config.append(f"max_client_conn = {max_conn}\n")
            new_config.append(f"default_pool_size = {default_pool_size}\n")
            new_config.append(f"min_pool_size = {min_pool_size}\n")
            new_config.append(f"reserve_pool_size = {reserve_pool_size}\n")
            new_config.append("pool_mode = session\n")
            new_config.append("server_reset_query = DISCARD ALL\n")
            new_config.append("server_check_delay = 30\n")
            new_config.append("server_check_query = SELECT 1\n")
            new_config.append("server_lifetime = 3600\n")
            new_config.append("server_idle_timeout = 600\n")
            new_config.append("query_wait_timeout = 120\n")
            new_config.append(f"tcp_keepalive = {tcp_keepalive}\n")
            new_config.append(f"tcp_keepidle = {tcp_keepidle}\n")
            new_config.append(f"tcp_keepintvl = {tcp_keepintvl}\n")
            new_config.append(f"tcp_keepcnt = {tcp_keepcnt}\n")
            new_config.append("listen_addr = *\n")
        
        # 写入新配置
        try:
            with open(self.config_file, 'w') as f:
                f.writelines(new_config)
            print_success("配置文件更新成功")
        except Exception as e:
            print_error(f"写入配置文件失败: {e}")
            print_info(f"尝试恢复备份...")
            try:
                shutil.copy2(backup_file, self.config_file)
                print_success("已恢复原配置")
            except Exception as e2:
                print_error(f"恢复备份失败: {e2}")
            return
        
        # 提示重载配置
        if confirm("是否立即重载PgBouncer配置?"):
            self.reload()
    
    def add_database(self):
        """添加数据库连接配置"""
        print_title("添加数据库连接配置")
        
        if not self._is_installed():
            print_error("PgBouncer未安装，请先安装")
            return
            
        if not os.path.exists(self.config_file):
            print_error(f"配置文件 {self.config_file} 不存在，请先安装PgBouncer")
            return
            
        # 获取数据库连接信息
        db_name = get_input("请输入数据库名称(pgbouncer中使用的名称)")
        db_host = get_input("请输入数据库主机地址", "localhost")
        db_port = get_input("请输入数据库端口", "5432")
        db_user = get_input("请输入数据库用户名", "postgres")
        db_password = get_input("请输入数据库密码(将添加到userlist.txt)")
        
        # 添加数据库配置
        try:
            # 读取当前配置
            with open(self.config_file, 'r') as f:
                config_lines = f.readlines()
            
            # 查找[databases]部分
            new_config = []
            in_databases_section = False
            databases_section_exists = False
            
            for line in config_lines:
                # 识别databases部分
                if line.strip() == "[databases]":
                    in_databases_section = True
                    databases_section_exists = True
                    new_config.append(line)
                    continue
                
                # 处理新部分开始
                if line.strip().startswith("[") and line.strip() != "[databases]":
                    if in_databases_section:
                        # 在离开databases部分前添加新条目
                        new_config.append(f"{db_name} = host={db_host} port={db_port} user={db_user}\n")
                    in_databases_section = False
                
                new_config.append(line)
            
            # 如果已经处理完所有行但仍在databases部分，添加新条目
            if in_databases_section:
                new_config.append(f"{db_name} = host={db_host} port={db_port} user={db_user}\n")
            
            # 如果没有找到databases部分，则添加
            if not databases_section_exists:
                new_config.insert(0, "[databases]\n")
                new_config.insert(1, f"{db_name} = host={db_host} port={db_port} user={db_user}\n\n")
            
            # 写入新配置
            with open(self.config_file, 'w') as f:
                f.writelines(new_config)
            
            print_success(f"数据库 {db_name} 配置添加成功")
            
            # 添加用户认证信息
            self._add_user_to_userlist(db_user, db_password)
            
        except Exception as e:
            print_error(f"添加数据库配置失败: {e}")
            return
        
        # 提示重载配置
        if confirm("是否立即重载PgBouncer配置?"):
            self.reload()
    
    def _add_user_to_userlist(self, username: str, password: str):
        """将用户添加到userlist.txt"""
        try:
            # 检查用户列表文件是否存在
            if not os.path.exists(self.userlist_file):
                # 创建空的用户列表文件
                self._create_empty_userlist()
            
            # 检查用户是否已存在
            user_exists = False
            with open(self.userlist_file, 'r') as f:
                for line in f:
                    if line.strip().startswith(f'"{username}"'):
                        user_exists = True
                        break
            
            # 计算密码MD5哈希
            md5_password = ""
            if password:
                # 使用pg_md5工具计算密码MD5
                if command_exists("pg_md5"):
                    try:
                        result = run_command(["pg_md5", password], capture_output=True)
                        md5_password = result.stdout.strip()
                    except Exception:
                        # 如果pg_md5失败，使用明文密码
                        md5_password = password
                else:
                    # 如果pg_md5不存在，使用明文密码
                    print_warning("未找到pg_md5工具，将使用明文密码")
                    md5_password = password
            
            # 添加或更新用户
            if user_exists:
                # 读取并更新文件
                with open(self.userlist_file, 'r') as f:
                    lines = f.readlines()
                
                with open(self.userlist_file, 'w') as f:
                    for line in lines:
                        if line.strip().startswith(f'"{username}"'):
                            f.write(f'"{username}" "{md5_password}"\n')
                        else:
                            f.write(line)
                print_success(f"更新用户 {username} 密码成功")
            else:
                # 添加新用户
                with open(self.userlist_file, 'a') as f:
                    f.write(f'"{username}" "{md5_password}"\n')
                print_success(f"添加用户 {username} 成功")
                
        except Exception as e:
            print_error(f"更新用户列表失败: {e}")

def main():
    """模块主函数，用于独立测试"""
    manager = PgBouncerManager()
    
    # 显示菜单
    print_title("PgBouncer管理工具")
    print("1. 安装PgBouncer")
    print("2. 卸载PgBouncer")
    print("3. 启动PgBouncer")
    print("4. 停止PgBouncer")
    print("5. 重启PgBouncer")
    print("6. 查看PgBouncer状态")
    print("7. 显示配置文件路径")
    print("8. 优化PgBouncer配置")
    print("9. 添加数据库连接配置")
    print("0. 退出")
    
    choice = input("\n请选择操作: ")
    
    if choice == "1":
        manager.install()
    elif choice == "2":
        manager.uninstall()
    elif choice == "3":
        manager.start()
    elif choice == "4":
        manager.stop()
    elif choice == "5":
        manager.restart()
    elif choice == "6":
        manager.status()
    elif choice == "7":
        manager.show_config()
    elif choice == "8":
        manager.optimize_config()
    elif choice == "9":
        manager.add_database()
    elif choice == "0":
        return
    else:
        print_error("无效的选择")
    
if __name__ == "__main__":
    main() 